# pip install -qU wikipedia
# -q 或 --quiet 静默模式安装（减少输出信息，不显示详细的安装过程） ; -U  或 --upgrade 升级到最新版本（如果已安装旧版本，会升级到最新版）
from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
from pydantic import BaseModel, Field

# from wikipedia import wikipedia
# wikipedia.set_lang("en")  # 设置为英文，通常更稳定
# wikipedia.set_rate_limiting(True)  # 启用速率限制


api_wrapper = WikipediaAPIWrapper(
    top_k_results=1,  # 返回前1个结果
    doc_content_chars_max=1000,  # 内容长度最多100
)

wiki_tool = WikipediaQueryRun(api_wrapper=api_wrapper)
print(f"wike_tool工具方法名称：{wiki_tool.name}")
print(f"wike_tool工具方法描述：{wiki_tool.description}")
print(f"wike_tool工具方法参数：{wiki_tool.args}")
print(f"wike_tool工具方法return_direct(控制是否绕过LLM直接返回结果)：{wiki_tool.return_direct}")
print(wiki_tool.invoke({"query": "china "}))


### 对维基百科自己再包一层 做一些拓展
class WikiInputs(BaseModel):
    """维基百科工具输入"""
    query: str = Field("查询维基百科内容，最多返回三个结果")

custom_wiki_tool = WikipediaQueryRun(
    name="wike_tool",
    description="维基百科查询",
    args_schema=WikiInputs,  # 自定义输入参数
    api_wrapper=api_wrapper,
    # 如果return_direct 设置为True，工具会直接返回查询结果，例如一个字符串或者一个简单的数据结构
    # 如果return_direct 设置为False, 工具可能会返回一个更复杂的响应对象，其中包含更多的元数据或结构化信息
    return_direct=True
)
print(f"custom_wiki_tool工具方法名称：{custom_wiki_tool.name}")
print(f"custom_wiki_tool工具方法描述：{custom_wiki_tool.description}")
print(f"custom_wiki_tool工具方法参数：{custom_wiki_tool.args}")
print(f"custom_wiki_tool工具方法return_direct(控制是否绕过LLM直接返回结果)：{custom_wiki_tool.return_direct}")
print(custom_wiki_tool.invoke({"query": "李雷 "}))